package com.yhzdys.ocr.service;

import ai.djl.Device;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Rectangle;
import ai.djl.ndarray.NDManager;
import ai.djl.onnxruntime.engine.OrtEngine;
import ai.djl.onnxruntime.engine.OrtNDManager;
import ai.djl.opencv.OpenCVImageFactory;
import ai.djl.pytorch.engine.PtNDManager;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.TranslateException;
import com.yhzdys.ocr.service.support.OcrExecThread;
import com.yhzdys.ocr.service.support.OcrUtils;
import com.yhzdys.ocr.service.support.PpWordDetectionTranslator;
import com.yhzdys.ocr.service.support.PpWordRecognitionTranslator;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.io.ByteArrayInputStream;
import java.lang.reflect.Field;
import java.net.URL;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

@Slf4j
@Service
public class OcrService {

    private static final List<NDManager> ND_MANAGERS; // for debug

    static {
        System.setProperty("ai.djl.pytorch.use_mkldnn", "false");
        System.setProperty("ai.djl.pytorch.cudnn_benchmark", "false");
        System.setProperty("ai.djl.pytorch.graph_optimizer", "false");
        System.setProperty("ai.djl.pytorch.num_interop_threads", "1");
        System.setProperty("ai.djl.pytorch.num_threads", "1");

        List<NDManager> ndManagers = new LinkedList<>();
        try {
            // PtNDManager
            Field ptSystemManagerField = PtNDManager.class.getDeclaredField("SYSTEM_MANAGER");
            ptSystemManagerField.setAccessible(true);
            PtNDManager ptNDManager = (PtNDManager) ptSystemManagerField.get(null);
            ptSystemManagerField.setAccessible(false);
            ndManagers.add(ptNDManager);
            // OrtNDManager
            Field ortSystemManagerField = OrtNDManager.class.getDeclaredField("SYSTEM_MANAGER");
            ortSystemManagerField.setAccessible(true);
            OrtNDManager ortNDManager = (OrtNDManager) ortSystemManagerField.get(null);
            ortSystemManagerField.setAccessible(false);
            ndManagers.add(ortNDManager);
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
        ND_MANAGERS = Collections.unmodifiableList(new ArrayList<>(ndManagers));
    }

    private final ZooModel<Image, DetectedObjects> detectionModel;
    private final ZooModel<Image, String> recognitionModel;
    /**
     * 异步执行OCR线程池
     */
    private final ThreadPoolExecutor executor;

    public OcrService() throws Exception {
        this.detectionModel = ModelZoo.loadModel(OcrService.detectionCriteria());
        this.recognitionModel = ModelZoo.loadModel(OcrService.recognitionCriteria());

        ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(
                2, 2,
                0L, TimeUnit.NANOSECONDS,
                new ArrayBlockingQueue<>(128),
                /* 线程池的线程，绑定固定的Predictor */
                runnable -> new OcrExecThread(runnable, detectionModel, recognitionModel),
                (runnable, executor) -> {
                    throw new RuntimeException("OCR任务队列已满，请稍后再试");
                }
        );
        threadPoolExecutor.allowCoreThreadTimeOut(false);
        this.executor = threadPoolExecutor;
    }

    private static Criteria<Image, DetectedObjects> detectionCriteria() throws Exception {
        return Criteria.builder()
                .optDevice(Device.cpu())
                .optEngine(OrtEngine.ENGINE_NAME)
                .optModelName("inference")
                .optModelPath(OcrService.loadModel("models/ch_PP-OCRv3_det_infer_onnx.zip"))
                .setTypes(Image.class, DetectedObjects.class)
                .optTranslator(new PpWordDetectionTranslator())
                .build();
    }

    private static Criteria<Image, String> recognitionCriteria() throws Exception {
        return Criteria.builder()
                .optDevice(Device.cpu())
                .optEngine(OrtEngine.ENGINE_NAME)
                .optModelName("inference")
                .optModelPath(OcrService.loadModel("models/ch_PP-OCRv3_rec_infer_onnx.zip"))
                .setTypes(Image.class, String.class)
                .optTranslator(new PpWordRecognitionTranslator())
                .build();
    }

    private static Path loadModel(String resource) throws Exception {
        URL url = OcrService.class.getClassLoader().getResource(resource);
        if (url == null) {
            throw new NullPointerException("Resource '" + resource + "' is null");
        }
        return Paths.get(url.toURI());
    }

    public Map<String, List<String>> recognize(List<String> urls) {
        if (CollectionUtils.isEmpty(urls)) {
            return Collections.emptyMap();
        }
        // map保存ocr识别结果
        Map<String, List<String>> ocrResultMap = new ConcurrentHashMap<>(urls.size(), 1.0F);
        CountDownLatch countDownLatch = new CountDownLatch(urls.size());

        for (String url : urls) {
            // 提交图片url给到线程池，并将识别结果存入ocrResultMap
            this.recognize(url, result -> {
                ocrResultMap.put(url, result);
                countDownLatch.countDown();
            });
        }

        boolean completed = false;
        try {
            completed = countDownLatch.await(120L, TimeUnit.SECONDS);
        } catch (Throwable ignored) {
        }
        if (completed) {
            return ocrResultMap;
        }
        // 补充下没有识别完的结果
        urls.forEach(url -> ocrResultMap.putIfAbsent(url, Collections.emptyList()));
        return ocrResultMap;
    }

    private void recognize(String url, Consumer<List<String>> consumer) {
        if (url == null || url.isEmpty()) {
            consumer.accept(Collections.emptyList());
            return;
        }
        URL finalUrl;
        try {
            finalUrl = new URL(url);
        } catch (Throwable t) {
            consumer.accept(Collections.emptyList());
            return;
        }
        try {
            // 提交到线程池，执行ocr
            executor.execute(() -> {
                List<String> result = this.recognize(finalUrl);
                // 处理ocr结果
                consumer.accept(result);
            });
        } catch (Throwable t) {
            log.error(t.getMessage(), t);
            consumer.accept(Collections.emptyList());
        }
    }

    private List<String> recognize(URL url) {
        Image image = null;
        try {
            byte[] bytes = IOUtils.toByteArray(url);
            image = OpenCVImageFactory.getInstance().fromInputStream(new ByteArrayInputStream(bytes));

            // 获取线程绑定的Predictor
            OcrExecThread currentThread = (OcrExecThread) Thread.currentThread();
            Predictor<Image, DetectedObjects> detector = currentThread.getDetector();
            Predictor<Image, String> recognizer = currentThread.getRecognizer();

            // 开始识别
            DetectedObjects detections = this.predict(image, detector, recognizer);

            List<String> result = detections.getClassNames();
            return CollectionUtils.isEmpty(result) ? Collections.emptyList() : result;
        } catch (Throwable t) {
            log.error(t.getMessage(), t);
            return Collections.emptyList();
        } finally {
            OcrUtils.releaseImage(image);
        }
    }

    public DetectedObjects predict(Image image, Predictor<Image, DetectedObjects> detector, Predictor<Image, String> recognizer) throws TranslateException {
        DetectedObjects detections = detector.predict(image);
        List<DetectedObjects.DetectedObject> boxes = detections.items();
        List<String> names = new LinkedList<>();
        List<Double> probabilities = new LinkedList<>();
        List<BoundingBox> boundingBoxes = new LinkedList<>();

        for (DetectedObjects.DetectedObject box : boxes) {
            Image subImg = null;
            try {
                subImg = this.getSubImage(image, box.getBoundingBox());
                String text = recognizer.predict(subImg);
                names.add(text);
                probabilities.add(-1.0);
                boundingBoxes.add(box.getBoundingBox());
            } finally {
                OcrUtils.releaseImage(subImg);
            }
        }
        return new DetectedObjects(names, probabilities, boundingBoxes);
    }

    private Image getSubImage(Image img, BoundingBox box) {
        Rectangle rect = box.getBounds();
        double[] extended = this.extendRect(rect.getX(), rect.getY(), rect.getWidth(), rect.getHeight());
        int width = img.getWidth();
        int height = img.getHeight();
        int[] recovered = {(int) (extended[0] * width), (int) (extended[1] * height), (int) (extended[2] * width), (int) (extended[3] * height)};
        return img.getSubImage(recovered[0], recovered[1], recovered[2], recovered[3]);
    }

    private double[] extendRect(double xMin, double yMin, double width, double height) {
        double centerX = xMin + width / 2;
        double centerY = yMin + height / 2;
        if (width > height) {
            width = width + height * 2.0;
            height *= 2.0;
        } else {
            height = height + width * 2.0;
            width *= 2.0;
        }
        double newX = centerX - width / 2 < 0 ? 0 : centerX - width / 2;
        double newY = centerY - height / 2 < 0 ? 0 : centerY - height / 2;
        double newWidth = newX + width > 1 ? 1 - newX : width;
        double newHeight = newY + height > 1 ? 1 - newY : height;
        return new double[]{newX, newY, newWidth, newHeight};
    }
}
